import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import transforms
import os
import unet
import data_handle
import numpy as np

from tqdm import tqdm
root = r"E:\YOLO\VOCdevkit\VOC2012"
weight_path = "./weight"
trans = transforms.Compose(
    [transforms.ToTensor()]
)



if __name__ == '__main__':

    model = unet.MainNet().cuda()
    if not os.path.exists(weight_path):
        os.mkdir(weight_path)
    dataloader = DataLoader(data_handle.UNET_dataloader(root),batch_size=3,shuffle=True)
    opt = torch.optim.Adam(model.parameters())
    loss_FC = nn.BCELoss()
    cls_pixel_list = torch.tensor(data_handle.cls_list)

    for epoch in range(100000):
        sum_loss= 0
        for i,(data,tag) in enumerate(tqdm(dataloader)):
            data = data.cuda()
            tag = tag.cuda()


            out = model(data)

            loss = loss_FC(out,tag)
            opt.zero_grad()
            loss.backward()
            opt.step()
            sum_loss += loss.cpu().detach().item()
        x = data[0].cpu()
        # 把标签还原成原图
        x_ = torch.argmax(tag[0], dim=0, keepdim=True)

        y = torch.argmax(out[0], dim=0, keepdim=True)

        tag_img = cls_pixel_list[x_][0].cpu().numpy()
        out_img = cls_pixel_list[y][0].cpu().numpy()

        tag_img = trans(tag_img)
        out_img = trans(out_img)
        z = torch.cat((x, tag_img, out_img), 2)
        img_save = transforms.ToPILImage()(z.cpu())
        img_save.save('image/{}.png'.format(epoch))
        print(f"epoch:{epoch},loss:{sum_loss/len(dataloader)}")
        torch.save(model.state_dict(),f'{weight_path}/{epoch}.pt')

